#pragma once
#include "common_unit.h"
#include <queue>
#include <map>
#include <mutex>

#include "memory_unit.h"
#define TIME_OUT_SECOND 5

class QueueUnit {
public:
    virtual ~QueueUnit() = default;

    virtual bool push(void* ptr) = 0;
    virtual bool pop(void * & ptr) = 0;
};

/* 用来分配内存的queue */
class MemQueueUnit : public QueueUnit {
public:
    MemQueueUnit(int MaxSize, std::vector<std::shared_ptr<MemoryNode> > & mem_node_list);

    bool push(void* ptr) override;
    bool pop(void * & ptr) override;

    /* 表示该mem queue作为某个设备传输前的最后一个stage，需要DMA数据传输到异构设备 */
    bool TriggerDmaSend();

    /* 表示该mem queue作为某个设备传输后第一个stage，需要接受异构设备DMA传输的数据*/
    bool TriggerDmaRecv();

    /* 初始化内存为某个值，方便调试 */
    void Mem_Init();

    //从mempool中补充内存指针到mem queue中
    void supply_from_mempool();

    /* 为MemoryPool 分配内存， 顺道初始化queue，全部塞满 */
    void mem_assign();
    size_t get_total_mem_size();
    size_t get_mem_size();

public:
    std::vector<void*> MemoryPool;
    std::queue<void*> Queue;
    std::mutex mtx;

    size_t Mem_element_size = 0; // every element size in memory pool
    int MaxQueueSize;

    stage_type_t stage_type; // if is DMA mem stage
    void * SendFlag;
    void * RecieveFlag;
    
    /* 在DMA中分配的task_id, Mem的有可能是个list，因为DMA有大小限制，mem可能超过这个大小。*/
    std::vector<int> Mem_task_id;
    int SendFlag_id;
    int RecvFlag_id;
};
/* 用来传递流水线的Queue */
class PipeQueueUnit : public QueueUnit {
public:
    PipeQueueUnit(int MaxSize,
                  std::shared_ptr<MemQueueUnit> FriendMemQueue,
                  bool if_first_stage,
                  bool if_last_stage);

    bool push(void* ptr) override;
    bool pop(void * & ptr) override;

private:
    int MaxQueueSize;
    stage_type_t stage_type;
    bool if_first_stage;
    bool if_last_stage;

    std::queue<void*> Queue;
    std::mutex mtx;

    std::shared_ptr<MemQueueUnit> FriendMemQueue;
};

/* function code implement */

/* MemQueue functions */

MemQueueUnit::MemQueueUnit(int MaxSize, std::vector<std::shared_ptr<MemoryNode> > & mem_node_list) {
    this->Mem_element_size = 0;
    this->MaxQueueSize = MaxSize;
    this->stage_type = stage_type_t::NORMAL_STAGE;

    for (auto mem_node : mem_node_list) {
        this->Mem_element_size += mem_node->total_size;
    }
}

size_t MemQueueUnit::get_total_mem_size() {
    return this->Mem_element_size * this->MaxQueueSize + sizeof(int64_t) + sizeof(int64_t);
}

size_t MemQueueUnit::get_mem_size() {
    return this->Mem_element_size * this->MaxQueueSize;
}

void MemQueueUnit::mem_assign() {
    // 确保是最后开启内存（等其他所有MemQueue+Stage都完成mem-assign再来
    MemoryUnit & memoryunit = MemoryUnit::Instance();
    this->MemoryPool.push_back(memoryunit.Malloc(this->get_total_mem_size()) );
    for (int i = 1; i < this->MaxQueueSize; ++i) {
        this->MemoryPool.push_back(MemoryPool[i-1] + Mem_element_size);
    }

    this->supply_from_mempool();

    this->SendFlag = this->MemoryPool[(int)MemoryPool.size() - 1] + Mem_element_size;
    this->RecieveFlag = this->SendFlag + sizeof(int64_t);

    *reinterpret_cast<volatile int64_t*>(this->RecieveFlag) = 0;
    *reinterpret_cast<volatile int64_t*>(this->SendFlag)    = 0;
}

bool MemQueueUnit::push(void* ptr) {
    std::lock_guard<std::mutex> lock(this->mtx);
    if ( (int)Queue.size() == MaxQueueSize) return false;

    if (this->stage_type == stage_type_t::RECV_STAGE) { // 如果是DMA阶段，不回收，一次性DMA,FIXME: 只有是first_stage才不回收
        return true;
    }

    Queue.push(ptr);
    return true;
}

bool MemQueueUnit::pop(void * & ptr) {
    std::lock_guard<std::mutex> lock(this->mtx);
    if (Queue.empty()) return false;
    ptr = Queue.front();
    Queue.pop();
    return true;
}

void MemQueueUnit::Mem_Init() {
    for (int i = 0; i < (int)MemoryPool.size(); ++i) {
        CommonUtils::mem_init( (int64_t*)MemoryPool[i], this->Mem_element_size);
    }
}

void MemQueueUnit::supply_from_mempool(){
    std::lock_guard<std::mutex> lock(this->mtx);
    if (!Queue.empty()) {LOG(INFO) << "Logic error";}
    for (int i = 0; i < this->MaxQueueSize; ++i) {
        // TODO: 如果是RECV-STAGE，不应该push
        this->Queue.push(MemoryPool[i]);
    } 
}

bool MemQueueUnit::TriggerDmaSend(){
    // 等recieve flag接到对侧发送的消息，准备完毕
    int useconds = 0;
    int step_us = 100;
    while ( *reinterpret_cast<volatile int64_t*>(this->RecieveFlag) == 0) {
        usleep(step_us);
        useconds += step_us;
        if (useconds > TIME_OUT_SECOND * 1e6) {
            LOG(WARNING) << "brpc connect overtime";
            return false;
        }
    }

    *reinterpret_cast<volatile int64_t*>(this->RecieveFlag) = 0;
    
    BrpcUnit & brpcunit = BrpcUnit::GetInstance();
    brpcunit.task_dispatch_mem (this->Mem_task_id);

    *reinterpret_cast<volatile int64_t*>(this->SendFlag)    = 1;
    brpcunit.task_dispatch_mem (std::vector<int>{this->SendFlag_id});

    this->supply_from_mempool();

    return true;
}

bool MemQueueUnit::TriggerDmaRecv(){
    // No extra mem ptr can use 
    BrpcUnit & brpcunit = BrpcUnit::GetInstance();
    // trigger DMA recieveFlag
    *reinterpret_cast<volatile int64_t*>(this->RecieveFlag) = 1;
    brpcunit.task_dispatch_mem (std::vector<int>{this->RecvFlag_id});

    // poll wait for send already
    int useconds = 0;
    int step_us = 100;
    while ( *reinterpret_cast<volatile int64_t*>(this->SendFlag) == 0) {
        usleep(step_us);
        useconds += step_us;
        if (useconds > TIME_OUT_SECOND * 1e6) {
            LOG(WARNING) << "brpc connect overtime";
            return false;
        }
    }
    *reinterpret_cast<volatile int64_t*>(this->SendFlag) = 0;

    this->supply_from_mempool();

    return true;
    //
}

/* PipeQueue function implement */

PipeQueueUnit::PipeQueueUnit(int MaxSize,
                  std::shared_ptr<MemQueueUnit> FriendMemQueue,
                  bool if_first_stage,
                  bool if_last_stage)
{
    this->FriendMemQueue = FriendMemQueue;
    this->stage_type = FriendMemQueue->stage_type;
    this->if_first_stage = if_first_stage;
    this->if_last_stage  = if_last_stage;
    this->MaxQueueSize   = MaxSize;                  
}

bool PipeQueueUnit::push(void* ptr) {
    std::lock_guard<std::mutex> lock(this->mtx);

    if (this->if_last_stage) {
        // directly use FriendMemQueue for recycle
        if (this->stage_type == stage_type_t::NORMAL_STAGE) {
            return this->FriendMemQueue->push(ptr);
        }
        else { // 如果是DMA，又是触发push，则必然是dma-send
            if ( (int)Queue.size() < MaxQueueSize) Queue.push(ptr);
            if ( (int)Queue.size() == MaxQueueSize) { // 可以DMA了
                bool return_value = (this->FriendMemQueue->TriggerDmaSend());
                if (return_value == false) return false;

                while (!Queue.empty()) { // DMA结束，继续流水线，内存全部存回 MemQueue
                    Queue.pop();
                }
            }
            return true;
        }
    }

    if ( (int)Queue.size() == MaxQueueSize) return false;
    Queue.push(ptr);
    return true;
}

bool PipeQueueUnit::pop(void * & ptr) {
    std::lock_guard<std::mutex> lock(this->mtx);

    if (this->if_first_stage) {
        // directly use FriendMemQueue
        if (this->stage_type == stage_type_t::NORMAL_STAGE) {
            return this->FriendMemQueue->pop(ptr);
        }
        else {
            bool ifempty = (this->FriendMemQueue->pop(ptr) == false);
            if (ifempty) { // 如果所有内存都分配出去了，请求发起 DMA
                bool return_value = this->FriendMemQueue->TriggerDmaRecv();
                if (return_value == false) return false;
                return this->FriendMemQueue->pop(ptr);
            }
            return true;
        }
    }

    if (Queue.empty()) return false;
    ptr = Queue.front();
    Queue.pop();
    return true;
}